"""
使用isaacgym程序包中的example/dof_controls修改。
控制walker2d的关节
"""

import math
from isaacgym import gymapi
from isaacgym import gymutil

# initialize gym
gym = gymapi.acquire_gym()

# parse arguments
args = gymutil.parse_arguments(description="Joint control Methods Example")

# create a simulator
sim_params = gymapi.SimParams()
sim_params.substeps = 2
sim_params.dt = 1.0 / 60.0

sim_params.physx.solver_type = 1
sim_params.physx.num_position_iterations = 4
sim_params.physx.num_velocity_iterations = 1

sim_params.physx.num_threads = args.num_threads
sim_params.physx.use_gpu = args.use_gpu

sim_params.use_gpu_pipeline = False
if args.use_gpu_pipeline:
    print("WARNING: Forcing CPU pipeline.")

sim = gym.create_sim(args.compute_device_id, args.graphics_device_id, args.physics_engine, sim_params)

if sim is None:
    print("*** Failed to create sim")
    quit()

# create viewer using the default camera properties
viewer = gym.create_viewer(sim, gymapi.CameraProperties())
if viewer is None:
    raise ValueError('*** Failed to create viewer')

# add ground plane
plane_params = gymapi.PlaneParams()
gym.add_ground(sim, gymapi.PlaneParams())

# set up the env grid
num_envs = 4
spacing = 1.5
env_lower = gymapi.Vec3(-spacing, 0.0, -spacing)
env_upper = gymapi.Vec3(spacing, 0.0, spacing)

# add cartpole urdf asset
asset_root = "../assets"
asset_file = "mjcf/walker2d_v2.xml"
# asset_file = "urdf/cartpole.urdf"

# Load asset with default control type of position for all joints
asset_options = gymapi.AssetOptions()
asset_options.fix_base_link = True
asset_options.default_dof_drive_mode = gymapi.DOF_MODE_POS
print("Loading asset '%s' from '%s'" % (asset_file, asset_root))
cartpole_asset = gym.load_asset(sim, asset_root, asset_file, asset_options)

# initial root pose for cartpole actors
initial_pose = gymapi.Transform()
initial_pose.p = gymapi.Vec3(0.0, 2.0, 0.0)
initial_pose.r = gymapi.Quat(-0.707107, 0.0, 0.0, 0.707107)

# Create environment 0
# Cart held steady using position target mode.
# Pole held at a 45 degree angle using position target mode.
env0 = gym.create_env(sim, env_lower, env_upper, 2)
cartpole0 = gym.create_actor(env0, cartpole_asset, initial_pose, 'cartpole', 0, 1)
# Configure DOF properties
props = gym.get_actor_dof_properties(env0, cartpole0)
# props["driveMode"] = (gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS, \
#                       gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS, \
#                       gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS)

props["driveMode"] = (gymapi.DOF_MODE_NONE, gymapi.DOF_MODE_NONE, gymapi.DOF_MODE_NONE, \
                      gymapi.DOF_MODE_VEL, gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS, \
                      gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS, gymapi.DOF_MODE_POS)


k1 = 50000.0
k1_0 = 0.0
k2 = 500.0
d1 = 1000.0
d1_0 = 0.0
d2 = 50.0
props["stiffness"] = (k1, k1, k1, \
                      0.0, k2, k2, \
                      k2, k2, k2)
props["damping"] = (d1, d1, d1, \
                    10*d2, d2, d2, \
                    d2, d2, d2)
gym.set_actor_dof_properties(env0, cartpole0, props)
# Set DOF drive targets

dof_names = ['ignorex', 'ignorez', 'ignorey', \
             'thigh_joint', 'leg_joint', 'foot_joint', \
             'thigh_left_joint', 'leg_left_joint', 'foot_left_joint']
angle = 0
flag = 0
for dof_name in dof_names:
    handle = gym.find_actor_dof_handle(env0, cartpole0, dof_name)
    if dof_name == 'thigh_joint':
        gym.set_dof_target_velocity(env0, handle, 30.0)

        # gym.set_dof_target_position(env0, handle, angle)
    else:
        gym.set_dof_target_position(env0, handle, 0)



# Look at the first env
cam_pos = gymapi.Vec3(8, 4, 1.5)
cam_target = gymapi.Vec3(0, 2, 1.5)
gym.viewer_camera_look_at(viewer, None, cam_pos, cam_target)

angle = 0
flag = 0

# Simulate
while not gym.query_viewer_has_closed(viewer):

    # step the physics
    gym.simulate(sim)
    gym.fetch_results(sim, True)

    # update the viewer
    gym.step_graphics(sim)
    gym.draw_viewer(viewer, sim, True)

    # Nothing to be done for env 0
    # dof_names = ['ignorex', 'ignorez', 'ignorey', \
    #              'thigh_joint', 'leg_joint', 'foot_joint', \
    #              'thigh_left_joint', 'leg_left_joint', 'foot_left_joint']
    #
    # for dof_name in dof_names:
    #     handle = gym.find_actor_dof_handle(env0, cartpole0, dof_name)
    #     if dof_name == 'thigh_joint':
    #         pos = gym.get_dof_position(env0, handle)
    #         if pos >= 10:
    #             gym.set_dof_target_velocity(env0, handle, -10.0)
    #         elif pos <= -10:
    #             gym.set_dof_target_velocity(env0, handle, 10.0)
    #
    #         # gym.set_dof_target_position(env0, handle, angle)
    #     else:
    #         gym.set_dof_target_position(env0, handle, 0)

    # Nothing to be done for env 1

    # # Update env 2: reverse cart target velocity when bounds reached
    # pos = gym.get_dof_position(env2, cart_dof_handle2)
    # if pos >= 0.5:
    #     gym.set_dof_target_velocity(env2, cart_dof_handle2, -1.0)
    # elif pos <= -0.5:
    #     gym.set_dof_target_velocity(env2, cart_dof_handle2, 1.0)
    #
    # # Update env 3: apply an effort to the pole to keep it upright
    # pos = gym.get_dof_position(env3, pole_dof_handle3)
    # gym.apply_dof_effort(env3, pole_dof_handle3, -pos * 50)

    # Wait for dt to elapse in real time.
    # This synchronizes the physics simulation with the rendering rate.
    gym.sync_frame_time(sim)

print('Done')

gym.destroy_viewer(viewer)
gym.destroy_sim(sim)
